import torch
from utils import set_seed, one_cycle, ModelEMA, EarlyStopping
from dataset import VibrationDataset
from model import SAFDNN
import numpy as np
import random
import os


if __name__ == "__main__":
    valid_dataset = VibrationDataset('SAFDNN/test_data_512.npy')
    valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=10, shuffle=True)
    model = SAFDNN(1, 4)
    if os.path.exists("best.pth"):
        model.load_state_dict(torch.load("best.pth"))
    correct = 0
    total = 0
    for idx, (valid_data, valid_label) in enumerate(valid_loader):

        output = model(valid_data.unsqueeze(1))
        pred = output.argmax(dim=1)
        correct += pred.eq(valid_label.view_as(pred)).sum().item()
        total += valid_label.size(0)
    print(correct / total)
